// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/common/selective_repeated_column_reader.h>
#include <pollux/dwio/parquet/reader/parquet_data.h>

namespace kumo::pollux::parquet {
    /// Comtainer for the lengths of a repeated reader where the lengths are
    /// pre-filled from repdefs.
    class RepeatedLengths {
    public:
        void setLengths(BufferPtr lengths) {
            lengths_ = std::move(lengths);
            nextLengthIndex_ = 0;
        }

        BufferPtr &lengths() {
            return lengths_;
        }

        int32_t nextLengthIndex() const {
            return nextLengthIndex_;
        }

        void readLengths(int32_t *lengths, int32_t numLengths) {
            POLLUX_CHECK_LE(
                nextLengthIndex_ + numLengths, lengths_->size() / sizeof(int32_t));
            memcpy(
                lengths,
                lengths_->as<int32_t>() + nextLengthIndex_,
                numLengths * sizeof(int32_t));
            nextLengthIndex_ += numLengths;
        }

    private:
        BufferPtr lengths_;
        int32_t nextLengthIndex_{0};
    };

    class MapColumnReader : public dwio::common::SelectiveMapColumnReader {
    public:
        MapColumnReader(
            const TypePtr &requestedType,
            const std::shared_ptr<const dwio::common::TypeWithId> &fileType,
            ParquetParams &params,
            common::ScanSpec &scanSpec);

        void prepareRead(
            vector_size_t offset,
            RowSet rows,
            const uint64_t *incomingNulls) {
            // The prepare is done by the topmost list/map/struct.
        }

        void seekToRowGroup(int64_t index) override;

        void enqueueRowGroup(uint32_t index, dwio::common::BufferedInput &input);

        void read(
            int64_t offset,
            const RowSet &rows,
            const uint64_t * /*incomingNulls*/) override;

        void setLengths(BufferPtr lengths) {
            lengths_.setLengths(lengths);
        }

        void readLengths(
            int32_t *lengths,
            int32_t numLengths,
            const uint64_t * /*nulls*/) override {
            lengths_.readLengths(lengths, numLengths);
        }

        /// Sets nulls and lengths of 'this' for the range of top level rows for which
        /// these have been decoded in 'leaf'.
        void setLengthsFromRepDefs(PageReader &leaf);

        /// advances 'this' to the end of the previously provided lengths/nulls. This
        /// is needed if lists are conditionally read from different structs that all
        /// end at different positions. Repeated children must use all lengths
        /// supplied before receiving new lengths.
        void skipUnreadLengths();

        void filterRowGroups(
            uint64_t rowGroupSize,
            const dwio::common::StatsContext &,
            dwio::common::FormatData::FilterRowGroupsResult &) const override;

    private:
        RepeatedLengths lengths_;
        RepeatedLengths keyLengths_;
        RepeatedLengths elementLengths_;
        LevelInfo levelInfo_;
    };

    class ListColumnReader : public dwio::common::SelectiveListColumnReader {
    public:
        ListColumnReader(
            const TypePtr &requestedType,
            const std::shared_ptr<const dwio::common::TypeWithId> &fileType,
            ParquetParams &params,
            common::ScanSpec &scanSpec);

        void prepareRead(
            vector_size_t offset,
            RowSet rows,
            const uint64_t *incomingNulls) {
            // The prepare is done by the topmost list/struct.
        }

        void seekToRowGroup(int64_t index) override;

        void enqueueRowGroup(uint32_t index, dwio::common::BufferedInput &input);

        void read(
            int64_t offset,
            const RowSet &rows,
            const uint64_t * /*incomingNulls*/) override;

        void setLengths(BufferPtr lengths) {
            lengths_.setLengths(lengths);
        }

        void readLengths(
            int32_t *lengths,
            int32_t numLengths,
            const uint64_t * /*nulls*/) override {
            lengths_.readLengths(lengths, numLengths);
        }

        /// Sets nulls and lengths of 'this' for the range of top level rows for which
        /// these have been decoded in 'leaf'.
        void setLengthsFromRepDefs(PageReader &leaf);

        /// advances 'this' to the end of the previously provided lengths/nulls. This
        /// is needed if lists are conditionally read from different structs that all
        /// end at different positions. Repeated children must use all lengths
        /// supplied before receiving new lengths.
        void skipUnreadLengths();

        void filterRowGroups(
            uint64_t rowGroupSize,
            const dwio::common::StatsContext &,
            dwio::common::FormatData::FilterRowGroupsResult &) const override;

    private:
        RepeatedLengths lengths_;
        LevelInfo levelInfo_;
    };

    /// Sets nulls and lengths for 'reader' and its children for the
    /// next 'numTop' top level rows. 'reader' must be a complex type
    /// reader. 'reader' may be inside structs but may not be inside a
    /// repeated reader. The topmost repeated reader ensures repdefs for
    /// all its children.
    void ensureRepDefs(dwio::common::SelectiveColumnReader &reader, int32_t numTop);
} // namespace kumo::pollux::parquet
